// See the file "COPYING" in the main distribution directory for copyright.

//
// An opinionated cluster-layout.zeek generator for single node deployments
// where the Zeek workers do not have ports allocated.
//
// The cluster-layout.zeek content is produced on stdout or written to the
// path given to the -o flag.
//
// Usage: zeek-cluster-layout-generator -L <loggers> -P <proxies> -W <workers>
// -a <listen_address> -p <port> -m <metrics_port> [-o outfile]
//
#include <unistd.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <ostream>
#include <string>
#include <vector>

namespace {

/**
 * Default is a small cluster for quick testing.
 */
struct ClusterLayoutOptions {
    int loggers = 1;
    int proxies = 1;
    int workers = 4;
    int port = 27760;
    std::string address = "127.0.0.1";
    int metrics_port = 9991;
    std::string metrics_address = "0.0.0.0";

    std::string NextPort() {
        auto result = std::to_string(port) + "/tcp";
        ++port;
        return result;
    }
    std::string NextMetricsPort() {
        auto result = std::to_string(metrics_port) + "/tcp";
        ++metrics_port;
        return result;
    }
};

class ClusterLayout {
public:
    ClusterLayout(ClusterLayoutOptions opts) {
        lines.emplace_back("# Auto-generated by zeek-cluster-layout-generator");
        lines.emplace_back("");
        lines.emplace_back("redef Cluster::manager_is_logger = " + std::string(opts.loggers > 0 ? "F" : "T") + ";");
        lines.emplace_back("");
        lines.emplace_back("redef Cluster::nodes += {");

        AddNode("manager", 0, "MANAGER", opts.address, opts.NextMetricsPort(), opts.NextPort());

        for ( int i = 1; i <= opts.loggers; i++ )
            AddNode("logger", i, "LOGGER", opts.address, opts.NextMetricsPort(), opts.NextPort());

        for ( int i = 1; i <= opts.proxies; i++ )
            AddNode("proxy", i, "PROXY", opts.address, opts.NextMetricsPort(), opts.NextPort());

        for ( int i = 1; i <= opts.workers; i++ )
            AddNode("worker", i, "WORKER", opts.address, opts.NextMetricsPort());

        lines.emplace_back("};");

        lines.emplace_back("");
        lines.emplace_back("@load base/frameworks/telemetry/options");
        lines.emplace_back("redef Telemetry::metrics_address = \"" + opts.metrics_address + "\";");
        lines.emplace_back("redef Telemetry::metrics_port = Cluster::local_node_metrics_port();");
    }

    void AddNode(const std::string& base_name, int node_idx, const std::string& type, const std::string& ip,
                 const std::string& metrics_port = "", const std::string& port = "") {
        std::string name = base_name;
        if ( node_idx > 0 )
            name += "-" + std::to_string(node_idx);
        auto s = "    [\"" + name + "\"] = [$node_type=Cluster::" + type + ", $ip=" + ip;

        if ( ! port.empty() )
            s += ", $p=" + port;

        if ( name != "manager" )
            s += ", $manager=\"manager\"";

        if ( ! metrics_port.empty() )
            s += ", $metrics_port=" + metrics_port;

        s += "],";
        lines.push_back(std::move(s));
    }

    void Print(std::ostream& out) {
        for ( const auto& l : lines )
            out << l << "\n";
    }

private:
    std::vector<std::string> lines;
};

void usage(const char* prog) {
    fprintf(stdout,
            "Usage: %s -L <loggers> -P <proxies> -W <workers> -a <address> -p <port> -m <metrics_port>"
            "-b <metrics_address> [-o <outfile>]\n",
            prog);
}

int checked_strtoul(int opt, const char* v) {
    const char* expected_end = v + std::strlen(v);
    char* end;
    auto r = std::strtoul(v, &end, 10);
    if ( end != expected_end ) {
        std::fprintf(stderr, "invalid numeric value for %c: %s\n", opt, v);
        std::exit(1);
    }

    return r;
}

} // namespace

int main(int argc, char* argv[]) {
    ClusterLayoutOptions opts;

    std::string out = "-";
    opterr = 0;

    while ( int c = getopt(argc, argv, "hL:P:W:a:p:m:b:o:") ) {
        if ( c < 0 )
            break;

        switch ( c ) {
            case 'L': opts.loggers = checked_strtoul(c, optarg); break;
            case 'P': opts.proxies = checked_strtoul(c, optarg); break;
            case 'W': opts.workers = checked_strtoul(c, optarg); break;
            case 'a': opts.address = optarg; break;
            case 'p': opts.port = checked_strtoul(c, optarg); break;
            case 'm': opts.metrics_port = checked_strtoul(c, optarg); break;
            case 'b': opts.metrics_address = optarg; break;
            case 'o': out = optarg; break;
            case 'h': usage(argv[0]); std::exit(0);
            case '?': std::fprintf(stderr, "unrecognized flag %c\n", optopt); std::exit(1);
        }
    }

    ClusterLayout layout(std::move(opts));

    // stdout print?
    if ( out == "-" ) {
        layout.Print(std::cout);
        return 0;
    }

    std::ofstream os(out, std::ios::trunc);
    if ( ! os.is_open() ) {
        std::fprintf(stderr, "failed to open %s\n", out.c_str());
        std::exit(1);
    }

    layout.Print(os);
    return 0;
}
